package com.freedom.rpc.thrift.common.utils;

import org.apache.thrift.protocol.TMultiplexedProtocol;
import org.apache.thrift.transport.TFramedTransport;

import com.freedom.rpc.thrift.common.extended.ExtendedTMultiplexedProtocol;
import com.freedom.rpc.thrift.common.extended.ServiceNameFieldHelper;
import com.freedom.rpc.thrift.common.extended.TransportFieldHelper;
import com.freedom.rpc.thrift.common.machine.MachineResource;
import com.freedom.rpc.thrift.common.socket.MyTSocket;
import com.freedom.rpc.thrift.common.socket.TSocketPool;

public class SocketUtils {
	private static final Logger logger = Logger.getLogger(SocketUtils.class);

	public static void setAlive(TMultiplexedProtocol tMultiProtocol, boolean alive, Exception e) {
		// 参数保护
		if (null == tMultiProtocol) {
			return;
		}
		try {
			TFramedTransport transport = (TFramedTransport) tMultiProtocol.getTransport();
			MyTSocket myTSocket = (MyTSocket) TransportFieldHelper.get(transport);
			myTSocket.setAlive(alive);
			logger.error("set socket to be not alive --- " + myTSocket.getSocket() + " due to :" + e);
		} catch (Exception e0) {
			logger.error(e0.toString());
		}
	}

	public static TMultiplexedProtocol getProtocol(String group, String service) throws Exception {
		TSocketPool pool = MachineResource.getInstance().getSocketPool(group, service);// pool
		if (null == pool) {
			throw new Exception("fail to fetch socket pool...");
		}
		TMultiplexedProtocol tMultiProtocol = pool.borrowObject();// tMultiProtocol
		if (null == tMultiProtocol) {
			throw new Exception("fail to fetch a socket from pool...");
		}
		logger.info("use " + TransportFieldHelper.get((TFramedTransport) tMultiProtocol.getTransport()).getSocket());
		// 注入本次的服务名
		ServiceNameFieldHelper.set(tMultiProtocol, service);
		// 保留socketPool句柄，为了后面释放使用
		ExtendedTMultiplexedProtocol extendedProtocol = (ExtendedTMultiplexedProtocol) tMultiProtocol;
		extendedProtocol.setPoolRef(pool);
		// 返回
		return tMultiProtocol;
	}

	public static void returnObject(TMultiplexedProtocol tMultiProtocol) {
		if (null == tMultiProtocol) {
			return;
		}
		ExtendedTMultiplexedProtocol extendedProtocol = (ExtendedTMultiplexedProtocol) tMultiProtocol;
		TSocketPool pool = extendedProtocol.getPoolRef();
		extendedProtocol.setPoolRef(null);// 拿出pool后，立即切断关联关系
		if (null != pool) {
			pool.returnObject(tMultiProtocol);
		}
	}
}
